package com.example.demo.util;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.*;
import java.io.File;
import java.io.IOException;
import java.util.*;
import java.util.List;

public class PHashUtil {

    private static final int SIZE = 32;
    private static final int DCT_SIZE = 8;
    private static final int SIMILARITY_THRESHOLD = 20; // 汉明距离阈值（<=10视为相似）

    // 生成pHash
    public static String getHash(String imagePath) throws IOException {
        BufferedImage img = ImageIO.read(new File(imagePath));

        // 1. 缩放32x32
        img = resize(img, SIZE, SIZE);

        // 2. 灰度化
        double[][] grayPixels = getGrayPixels(img);

        // 3. DCT变换
        double[][] dct = applyDCT(grayPixels);

        // 4. 计算均值（取8x8左上区域）
        double avg = getAvg(dct);

        // 5. 生成哈希
        return generateHash(dct, avg);
    }

    // 调整图片大小
    private static BufferedImage resize(BufferedImage src, int w, int h) {
        BufferedImage dest = new BufferedImage(w, h, BufferedImage.TYPE_INT_ARGB);
        Graphics2D g = dest.createGraphics();
        g.setRenderingHint(RenderingHints.KEY_INTERPOLATION,
                RenderingHints.VALUE_INTERPOLATION_BILINEAR);
        g.drawImage(src, 0, 0, w, h, null);
        g.dispose();
        return dest;
    }

    // 获取灰度像素矩阵
    private static double[][] getGrayPixels(BufferedImage img) {
        int width = img.getWidth();
        int height = img.getHeight();
        double[][] pixels = new double[height][width];

        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                int rgb = img.getRGB(x, y);
                int r = (rgb >> 16) & 0xFF;
                int g = (rgb >> 8) & 0xFF;
                int b = rgb & 0xFF;
                // 灰度公式
                pixels[y][x] = 0.299 * r + 0.587 * g + 0.114 * b;
            }
        }
        return pixels;
    }

    // DCT变换实现
    private static double[][] applyDCT(double[][] input) {
        int size = input.length;
        double[][] output = new double[size][size];

        for (int u = 0; u < size; u++) {
            for (int v = 0; v < size; v++) {
                double sum = 0.0;
                for (int i = 0; i < size; i++) {
                    for (int j = 0; j < size; j++) {
                        sum += input[i][j] *
                                Math.cos((2 * i + 1) * u * Math.PI / (2.0 * size)) *
                                Math.cos((2 * j + 1) * v * Math.PI / (2.0 * size));
                    }
                }
                double cu = u == 0 ? 1.0 / Math.sqrt(2) : 1.0;
                double cv = v == 0 ? 1.0 / Math.sqrt(2) : 1.0;
                output[u][v] = 0.25 * cu * cv * sum;
            }
        }
        return output;
    }

    // 计算DCT系数均值
    private static double getAvg(double[][] dct) {
        double sum = 0;
        for (int i = 0; i < DCT_SIZE; i++) {
            for (int j = 0; j < DCT_SIZE; j++) {
                sum += dct[i][j];
            }
        }
        return sum / (DCT_SIZE * DCT_SIZE);
    }

    // 生成二进制哈希
    private static String generateHash(double[][] dct, double avg) {
        StringBuilder hash = new StringBuilder();
        for (int i = 0; i < DCT_SIZE; i++) {
            for (int j = 0; j < DCT_SIZE; j++) {
                hash.append(dct[i][j] > avg ? "1" : "0");
            }
        }
        return hash.toString();
    }

    // 计算汉明距离
    public static int getDistance(String hash1, String hash2) {
        int distance = 0;
        for (int i = 0; i < hash1.length(); i++) {
            if (hash1.charAt(i) != hash2.charAt(i)) {
                distance++;
            }
        }
        return distance;
    }

    public static List<SimilarImage> findSimilarImages(String targetHash, File folder) {
        List<SimilarImage> result = new ArrayList<>();
        try {
            Arrays.stream(Objects.requireNonNull(folder.listFiles()))
                    .parallel()  // 并行处理提升性能
                    .filter(f -> f.isFile() && isImageFile(f))
                    .forEach(f -> {
                        try {
                            String currentHash = getHash(f.getAbsolutePath());
                            int distance = getDistance(targetHash, currentHash);

                            if (distance <= SIMILARITY_THRESHOLD) {
                                result.add(new SimilarImage(
                                        f.getAbsolutePath(),
                                        distance
                                ));
                            }
                        } catch (Exception e) {
                            System.err.println("无法读取文件: " + f.getName());
                        }
                    });
        } catch (Exception e) {
            throw new RuntimeException("图片处理失败", e);
        }
        // 按汉明距离排序
        Collections.sort(result);
        return result;
    }

    // 判断是否为图片文件
    private static boolean isImageFile(File file) {
        String name = file.getName().toLowerCase();
        return file.isFile() && (
                name.endsWith(".jpg") ||
                        name.endsWith(".jpeg") ||
                        name.endsWith(".png") ||
                        name.endsWith(".bmp")
        );
    }

    // 相似图片数据类
    static class SimilarImage implements Comparable<SimilarImage> {
        String filePath;
        int distance;

        SimilarImage(String path, int dist) {
            this.filePath = path;
            this.distance = dist;
        }

        @Override
        public int compareTo(SimilarImage o) {
            return Integer.compare(this.distance, o.distance);
        }
    }

    public static void main(String[] args) throws IOException {
//        String hash1 = getHash("D:/input/01-05.jpg");
//        String hash2 = getHash("D:/input/009_11.png");
//        int distance = getDistance(hash1, hash2);
//
//        System.out.println("Hash1: " + hash1);
//        System.out.println("Hash2: " + hash2);
//        System.out.println("Hamming Distance: " + distance);
//        System.out.println("Similarity: " + (100 - (distance * 100) / 64) + "%");

        String targetImage = "D:/input/01-05.jpg";
        String folderPath = "D:/input1";

        try {
            // 获取目标图片哈希
            String targetHash = getHash(targetImage);

            // 遍历文件夹获取相似图片
            List<SimilarImage> results = findSimilarImages(
                    targetHash,
                    new File(folderPath)
            );

            // 输出结果
            System.out.println("相似图片列表（按相似度降序）：");
            for (SimilarImage sim : results) {
                System.out.printf("[距离%d] %s => 相似度%d%%\n",
                        sim.distance,
                        sim.filePath,
                        100 - (sim.distance * 100) / 64
                );
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}